package com.sawyerqiu.agent;

import com.sun.tools.attach.VirtualMachine;

import java.io.IOException;
import java.lang.instrument.Instrumentation;
import java.lang.management.ManagementFactory;
import java.lang.reflect.Field;
import java.nio.file.Files;
import java.nio.file.Path;
import java.nio.file.Paths;
import java.util.Optional;
import java.util.regex.Pattern;

/**
 * @author : SawyerQiu
 * @date : 2024/7/31 18:19
 */
public class InstrumentUtil {

    private static final Pattern AGENT_JAR_NAME_PARTTERN = Pattern.compile("^my-agent-[0-9.]*(-SNAPSHOT){0,1}.jar$");

    private String agentJarPath;

    private Instrumentation instrumentation;

    public InstrumentUtil(String libPathName) throws IOException {
        Path libPath = Paths.get(libPathName);
        if (!Files.exists(libPath)) {
            System.out.println("无法启动热更，错误的libPath" + libPathName);
            return;
        }

        Optional<Path> find = Files.find(Paths.get(libPathName), 1,
                (path, fileAttr) -> AGENT_JAR_NAME_PARTTERN.matcher(path.getFileName().toString()).find())
                .findAny();

        if (!find.isPresent()) {
            System.out.println(String.format("无法启动热更，libPath[%s]下未找到[%s]", libPath, AGENT_JAR_NAME_PARTTERN));
            return;
        }

        Path path = find.get();
        this.agentJarPath = path.toAbsolutePath().toString();
        System.out.println(String.format("热更启动，找到包[%s]", agentJarPath));
    }

    public synchronized Instrumentation getInstrumentation() {
        if (instrumentation == null && agentJarPath != null) {
            VirtualMachine vm = null;
            try {
                String pid = ManagementFactory.getRuntimeMXBean().getName().split("@")[0];
                vm = VirtualMachine.attach(pid);
                vm.loadAgent(agentJarPath, null);
                instrumentation = doGetInstrumentation();

            } catch (Throwable e) {
                System.err.println("getInstrumentation fail:");
                e.printStackTrace();
            } finally {
                if (vm != null) {
                    try {
                        vm.detach();;
                    } catch (IOException e) {
                        e.printStackTrace();
                    }
                }
            }
        }
        return instrumentation;
    }

    private Instrumentation doGetInstrumentation() throws ClassNotFoundException, NoSuchFieldException, IllegalAccessException {
        Class<?> agentClass = Class.forName("com.sawyerqiu.agent.MyAgent");
        Field instrumentationField = agentClass.getDeclaredField("instrumentation");
        return (Instrumentation) instrumentationField.get(null);
    }
}
